import numpy as np
import keras
import secretary
from sklearn.linear_model import LinearRegression
from river.forest import ARFRegressor

import joblib
import os

from river.datasets.base import Dataset, MULTI_CLF

class CustomDataset(Dataset):
    
    def __init__(self, x_train_ml, y_train_ml):
        super().__init__(
            task=MULTI_CLF,
            n_features=16
        )
        self.x_train_ml = []
        for x_one in x_train_ml:
            x_buffer = {}
            for i, x in enumerate(x_one):
                x_buffer.update({i: x})
            self.x_train_ml.append(x_buffer)
        self.y_train_ml = y_train_ml
    
    def __iter__(self):
        for x, y in zip(self.x_train_ml, self.y_train_ml):
            yield x, y

def predict(usage_data, solar_data):
    """use neural network to predict the data

    Args:
        usage_data (np.array): 15 climate data
        solar_data (np.array): 5 climate data

    Returns:
        tuple: two predicted result, shape as [24]
    """    
    usage_data = usage_data.reshape(1, 24, 15)
    lstm_model = keras.models.load_model('lstm_model.keras')
    arfr_model = joblib.load('htr_model.joblib')
    secretary.log(f'lstm_in: {usage_data}', workbook='./programme.log')
    lstm_out = lstm_model.predict(usage_data)
    secretary.log(f'lstm_out: {lstm_out}', workbook='./programme.log')
    usage_data = usage_data.reshape(24, 15)
    print('[debug]',usage_data.shape)
    lstm_out = lstm_out.reshape(24, 1)
    print('[debug]',usage_data.shape, lstm_out.shape)
    lstm_out = np.concatenate((usage_data, lstm_out), axis=1)
    secretary.log(f'htr_in: {lstm_out}')
    usage = []
    for predict_pack in lstm_out:
        arfr_input = {}
        for i, x in enumerate(predict_pack):
            arfr_input.update({i: x})
        usage.append(arfr_model.predict_one(arfr_input)[0])
        print(f'[debug] predict input {arfr_input}, output {usage[-1]}')
    secretary.log(f'htr_out {usage}')
    
    lr_model = joblib.load('LinearRegression.joblib')
    solar_data = solar_data.reshape(-1, 5)
    secretary.log(f'lr_in: {solar_data}')
    solar = lr_model.predict(solar_data)
    solar = solar.reshape(-1)
    secretary.log(f'lr_out: {solar}')
    return usage, solar

def planner(storage, usage, solar):
    """calculate how much electricity do we need to pull during the night

    Args:
        storage (double): the current power storage
        usage (np.array): predict result
        solar (np.array): predict

    Returns:
        double: the volume of electricity
    """    
    if usage is np.array and solar is np.array:
        pull = 0
        for i in range(24):
            storage += (solar[i] - usage[i])
            if storage < 0:
                pull -= storage
                storage = 0
        return pull
    else:
        return None

def update(usage_data, consume, alpha=0.05):
    """update weights online

    Args:
        usage_data ([n,24,15]): recent usage data storaged in database
        consume ([n,24]): exact power consumption storaged in database
        alpha (float, optional): learning rate. Defaults to 0.05.
    """    
    lstm_model = keras.models.load_model('lstm_model.keras')
    old_weights = lstm_model.get_weights()
    usage_data, consume = usage_data.reshape(-1, 24, 15), consume.reshape(-1, 24)
    lstm_model.fit(usage_data, consume, batch_size=1, epochs=1, verbose=0)
    new_weights = lstm_model.get_weights()
    combined_weights = [(1 - alpha) * ow + alpha * nw for ow, nw in zip(old_weights, new_weights)]
    lstm_model.set_weights(combined_weights)
    lstm_model.save('lstm_model.keras')
    
    arfr_train = lstm_model.predict(usage_data)
    usage_data, arfr_train, consume = usage_data.reshape(-1, 15), arfr_train.reshape(-1, 1), consume.reshape(-1, 1)
    arfr_train = np.concatenate([usage_data, arfr_train], axis=1)
    
    arfr_model = joblib.load('htr_model.joblib')
    arfr_train = CustomDataset(arfr_train, consume)
    for x, y in arfr_train:
        arfr_model.learn_one(x, y)
    joblib.dump(arfr_model, 'htr_model.joblib')

if __name__ == '__main__':
    assert os.path.exists('lstm_model.keras'), "LSTM model file not found"
    assert os.path.exists('htr_model.joblib'), "Random Forest model file not found"
    assert os.path.exists('LinearRegression.joblib'), "LinearRegression model file not found"
    model = keras.models.load_model('lstm_model.keras')
    del model
